import torch
from MyData import MyDataset
from torch.utils.data import DataLoader
from net import Model
from transformers import BertTokenizer
# Transformers 4.50.3 已移除 optimization.AdamW 的显式导出，或需依赖 PyTorch 的 torch.optim.AdamW 实现‌
from torch.optim import AdamW

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
EPOCH = 30000

token = BertTokenizer.from_pretrained(
    r"../../my_model_cache/bert-base-chinese/models--bert-base-chinese/snapshots/c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f")


def collate_fn(data):
    sentes = [i[0] for i in data]
    label = [i[1] for i in data]
    # print(sentes)
    #编码
    data = token.batch_encode_plus(batch_text_or_text_pairs=sentes,
                                   truncation=True,
                                   padding="max_length",
                                   max_length=512,
                                   return_tensors="pt",
                                   return_length=True)
    input_ids = data["input_ids"]
    attention_mask = data["attention_mask"]
    token_type_ids = data["token_type_ids"]
    labels = torch.LongTensor(label)
    # print(input_ids,attention_mask,token_type_ids)
    return input_ids, attention_mask, token_type_ids, labels


#创建数据集
train_dataset = MyDataset("train")
val_dataset = MyDataset("validation")
#创建dataloader
train_laoder = DataLoader(dataset=train_dataset,
                          batch_size=100,
                          shuffle=True,
                          drop_last=True,
                          collate_fn=collate_fn)
val_loader = DataLoader(dataset=val_dataset,
                        batch_size=50,
                        shuffle=True,
                        drop_last=True,
                        collate_fn=collate_fn)

if __name__ == '__main__':
    #开始训练
    print(DEVICE)
    model = Model().to(DEVICE)
    optimizer = AdamW(model.parameters(), lr=5e-4)
    loss_func = torch.nn.CrossEntropyLoss()

    model.train()
    for epoch in range(EPOCH):
        sum_val_acc = 0
        sum_val_loss = 0
        #训练
        for i, (input_ids, attention_mask, token_type_ids, labels) in enumerate(train_laoder):
            # print(input_ids)
            input_ids, attention_mask, token_type_ids, labels = input_ids.to(DEVICE), attention_mask.to(
                DEVICE), token_type_ids.to(DEVICE), labels.to(DEVICE)
            out = model(input_ids, attention_mask, token_type_ids)

            loss = loss_func(out, labels)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if i % 5 == 0:
                out = out.argmax(dim=1)
                acc = (out == labels).sum().item() / len(labels)
                print(epoch, i, loss.item(), acc)
        # #模型验证
        model.eval()
        for i, (input_ids, attention_mask, token_type_ids, labels) in enumerate(val_loader):
            input_ids, attention_mask, token_type_ids, labels = input_ids.to(DEVICE), attention_mask.to(
                DEVICE), token_type_ids.to(DEVICE), labels.to(DEVICE)
            out = model(input_ids, attention_mask, token_type_ids)

            loss = loss_func(out, labels)

            out = out.argmax(dim=1)
            accuracy = (out == labels).sum().item() / len(labels)
            sum_val_loss = sum_val_loss + loss
            sum_val_acc = sum_val_acc + accuracy
        avg_val_loss = sum_val_loss / len(val_loader)
        avg_val_acc = sum_val_acc / len(val_loader)
        print(f"val==>epoch:{epoch},avg_val_loss:{avg_val_loss}, avg_val_acc:{avg_val_acc}")

        #if avg_val_acc > best_avg_val_acc:
        #   torch.save()
        #   best_avg_val_acc = avg_val_acc
        torch.save(model.state_dict(), f"params/{epoch}bert-weibo.pth")
        print(epoch, "参数保存成功！")
